#include "analyzer_common.h"

void bind_to_core(int core)
{
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(core, &cpuset);

    if (pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset) != 0)
    {
        perror("pthread_setaffinity_np error");
    }
}

int binarylowerbound(const analyzer_int64_t *arr, int len, analyzer_int64_t value)
{
    int left = 0;
    int right = len;
    while (left < right)
    {
        int mid = (left + right) / 2;
        // value <= arr[mid] ? (right = mid) : (left = mid + 1);
        value < arr[mid] ? (right = mid) : (left = mid + 1);
    }
    return left;
}

void exclusive_scan_1(analyzer_int64_t *input, int length)
{
    if (length == 0 || length == 1)
        return;

    analyzer_int64_t old_val, new_val;

    old_val = input[0];
    input[0] = 0;
    for (int i = 1; i < length; i++)
    {
        new_val = input[i];
        input[i] = old_val + input[i - 1];
        old_val = new_val;
    }
}

void exclusive_scan_3(unsigned int *input, int length)
{
    if (length == 0 || length == 1)
        return;

    unsigned int old_val, new_val;

    old_val = input[0];
    input[0] = 0;
    for (int i = 1; i < length; i++)
    {
        new_val = input[i];
        input[i] = old_val + input[i - 1];
        old_val = new_val;
    }
}

void analyzer_csc_to_bcsc(
    analyzer_exblock_idx n,
    analyzer_inblock_idx nb,
    
    analyzer_exblock_ptr* csc_pointer,
    analyzer_exblock_idx* csc_index,
    sparse_value_t* csc_value,

    analyzer_exblock_ptr** bcsc_struct_pointer,
    analyzer_exblock_idx** bcsc_struct_index,
    analyzer_exblock_ptr** bcsc_struct_nnzptr,
    analyzer_inblock_ptr*** bcsc_inblock_pointers,
    analyzer_inblock_idx*** bcsc_inblock_indeces,
    sparse_value_t*** bcsc_values
){
#define _PANGULU_SET_BVALUE_SIZE(size) ((csc_value)?(size):(0))

    char* block_csc = NULL;

    analyzer_int64_t block_length = ICEIL(n, nb);

    int nthreads = sysconf(_SC_NPROCESSORS_ONLN);
    #pragma omp parallel num_threads(nthreads)
    {
        nthreads = omp_get_num_threads();
        bind_to_core(omp_get_thread_num());
    }

    // printf("nthreads = %d\n", nthreads);

    analyzer_exblock_ptr *bcolptr = NULL;
    analyzer_int32_t *browidx = NULL;
    analyzer_int64_t *bnnzptr = NULL;

    analyzer_exblock_ptr nnz = csc_pointer[n];
    analyzer_int64_t bit_length = (block_length + 31) / 32;
    analyzer_exblock_ptr block_num = 0;
    analyzer_int64_t avg_nnz = ICEIL(nnz, nthreads);
    analyzer_int64_t *block_row_nnz_pt = (analyzer_int64_t *)malloc(sizeof(analyzer_int64_t) * (block_length + 1));
    for (int i = 0; i < block_length; i++)
    {
        block_row_nnz_pt[i] = csc_pointer[MIN(i * nb, n)];
    }
    block_row_nnz_pt[block_length] = csc_pointer[n];

    int *thread_pt = (int *)malloc(sizeof(int) * (nthreads + 1));
    thread_pt[0] = 0;
    for (int i = 1; i < nthreads + 1; i++)
    {
        thread_pt[i] = binarylowerbound(block_row_nnz_pt, block_length, avg_nnz * i);
    }
    free(block_row_nnz_pt);
    block_row_nnz_pt = NULL;

    bcolptr = (analyzer_exblock_ptr *)malloc(sizeof(analyzer_exblock_ptr) * (block_length + 1));
    memset(bcolptr, 0, sizeof(analyzer_exblock_ptr) * (block_length + 1));

    unsigned int *bit_array = (unsigned int *)malloc(sizeof(unsigned int) * bit_length * nthreads);

    #pragma omp parallel num_threads(nthreads)
    {
        int tid = omp_get_thread_num();
        unsigned int *tmp_bit = bit_array + bit_length * tid;

        for (int level = thread_pt[tid]; level < thread_pt[tid + 1]; level++)
        {
            memset(tmp_bit, 0, sizeof(unsigned int) * bit_length);

            int start_row = level * nb;
            int end_row = ((level + 1) * nb) < n ? ((level + 1) * nb) : n;

            for (int rid = start_row; rid < end_row; rid++)
            {
                for (analyzer_int64_t idx = csc_pointer[rid]; idx < csc_pointer[rid + 1]; idx++)
                {
                    analyzer_int32_t colidx = csc_index[idx];
                    analyzer_int32_t block_cid = colidx / nb;
                    setbit(tmp_bit[block_cid / 32], block_cid % 32);
                }
            }

            analyzer_int64_t tmp_blocknum = 0;
            for (int i = 0; i < bit_length; i++)
            {
                tmp_blocknum += __builtin_popcount(tmp_bit[i]);
            }

            bcolptr[level] = tmp_blocknum;
        }
    }
    exclusive_scan_1(bcolptr, block_length + 1);
    block_num = bcolptr[block_length];

    printf("block count : %lld\n", block_num);

    bnnzptr = (analyzer_int64_t *)malloc(sizeof(analyzer_int64_t) * (block_num + 1));
    memset(bnnzptr, 0, sizeof(analyzer_int64_t) * (block_num + 1));
    browidx = (analyzer_int32_t *)malloc(sizeof(analyzer_int32_t) * block_num);

    int *count_array = (int *)malloc(sizeof(int) * block_length * nthreads);
    #pragma omp parallel num_threads(nthreads)
    {
        int tid = omp_get_thread_num();
        unsigned int *tmp_bit = bit_array + bit_length * tid;
        int *tmp_count = count_array + block_length * tid;

        for (int level = thread_pt[tid]; level < thread_pt[tid + 1]; level++)
        {
            memset(tmp_bit, 0, sizeof(unsigned int) * bit_length);
            memset(tmp_count, 0, sizeof(int) * block_length);

            analyzer_int64_t *cur_block_nnz_pt = bnnzptr + bcolptr[level];
            analyzer_int32_t *cur_block_col_idx = browidx + bcolptr[level];

            int start_col = level * nb;
            int end_col = ((level + 1) * nb) < n ? ((level + 1) * nb) : n;

            for (int col = start_col; col < end_col; col++)
            {
                for (analyzer_int64_t idx = csc_pointer[col]; idx < csc_pointer[col + 1]; idx++)
                {
                    analyzer_int32_t row = csc_index[idx];
                    analyzer_int32_t brow = row / nb;
                    setbit(tmp_bit[brow / 32], brow % 32);
                    tmp_count[brow]++;
                }
            }

            analyzer_int64_t cnt = 0;
            for (int i = 0; i < block_length; i++)
            {
                if (getbit(tmp_bit[i / 32], i % 32))
                {
                    cur_block_nnz_pt[cnt] = tmp_count[i];
                    cur_block_col_idx[cnt] = i;
                    cnt++;
                }
            }
        }
    }
    free(bit_array);
    bit_array = NULL;
    free(count_array);
    count_array = NULL;
    exclusive_scan_1(bnnzptr, block_num + 1);

    // printf("nnz from bnnzptr = %lld\n", bnnzptr[block_num]);

    block_csc = malloc(
        sizeof(analyzer_inblock_ptr) * (nb + 1) * (block_num) +
        sizeof(analyzer_inblock_idx) * (nnz) + 
        _PANGULU_SET_BVALUE_SIZE(sizeof(sparse_value_t)) * (nnz)
    );
    memset(
        block_csc, 0, 
        sizeof(analyzer_inblock_ptr) * (nb + 1) * (block_num) +
        sizeof(analyzer_inblock_idx) * (nnz) + 
        _PANGULU_SET_BVALUE_SIZE(sizeof(sparse_value_t)) * (nnz)
    );

    #pragma omp parallel num_threads(nthreads)
    {
        int tid = omp_get_thread_num();
        int* tmp_count = malloc(sizeof(int) * block_length);

        for (int level = thread_pt[tid]; level < thread_pt[tid + 1]; level++)
        {
            memset(tmp_count, 0, sizeof(int) * block_length);
            analyzer_exblock_idx start_col = level * nb;
            analyzer_exblock_idx end_col = ((level + 1) * nb) < n ? ((level + 1) * nb) : n;
            for (analyzer_exblock_idx col = start_col, col_in_blc = 0; col < end_col; col++, col_in_blc++)
            {
                analyzer_exblock_ptr bidx = bcolptr[level];
                analyzer_exblock_idx brow = browidx[bidx];

                analyzer_inblock_ptr *cur_block_colptr = (analyzer_inblock_ptr *)(block_csc + sizeof(analyzer_inblock_ptr) * bidx * (nb + 1));
                analyzer_inblock_idx *cur_block_rowidx = (analyzer_inblock_idx *)(
                    block_csc
                    + sizeof(analyzer_inblock_ptr) * bcolptr[block_length] * (nb + 1) 
                    + sizeof(analyzer_inblock_idx) * bnnzptr[bidx]
                );
                sparse_value_t *cur_block_value = NULL;
                if(csc_value){
                    cur_block_value = (sparse_value_t *)(
                        block_csc
                        + sizeof(analyzer_inblock_ptr) * bcolptr[block_length] * (nb + 1) 
                        + sizeof(analyzer_inblock_idx) * bnnzptr[bcolptr[block_length]]
                        + sizeof(sparse_value_t) * bnnzptr[bidx]
                    );
                }

                for (analyzer_exblock_ptr idx = csc_pointer[col]; idx < csc_pointer[col + 1]; idx++)
                {
                    analyzer_exblock_idx row = csc_index[idx];
                    analyzer_exblock_idx new_brow = row / nb;
                    if(new_brow > brow){
                        while(new_brow != browidx[bidx]){
                            bidx++;
                        }
                        brow = new_brow;
                        cur_block_colptr = (analyzer_inblock_ptr *)(block_csc + sizeof(analyzer_inblock_ptr) * bidx * (nb + 1));
                        cur_block_rowidx = (analyzer_inblock_idx *)(
                            block_csc
                            + sizeof(analyzer_inblock_ptr) * bcolptr[block_length] * (nb + 1) 
                            + sizeof(analyzer_inblock_idx) * bnnzptr[bidx]
                        );
                        if(csc_value){
                            cur_block_value = (sparse_value_t *)(
                                block_csc
                                + sizeof(analyzer_inblock_ptr) * bcolptr[block_length] * (nb + 1) 
                                + sizeof(analyzer_inblock_idx) * bnnzptr[bcolptr[block_length]]
                                + sizeof(sparse_value_t) * bnnzptr[bidx]
                            );
                        }
                    }
                    cur_block_rowidx[tmp_count[brow]] = row % nb;
                    cur_block_colptr[col_in_blc]++;
                    if(csc_value){
                        cur_block_value[tmp_count[brow]] = csc_value[idx];
                    }
                    tmp_count[brow]++;
                }
            }
            for (analyzer_int64_t bidx = bcolptr[level]; bidx < bcolptr[level + 1]; bidx++)
            {
                analyzer_inblock_ptr *cur_block_colptr = (analyzer_inblock_ptr *)(block_csc + sizeof(analyzer_inblock_ptr) * bidx * (nb + 1));
                exclusive_scan_3(cur_block_colptr, nb + 1);
            }
        }

        free(tmp_count);
        tmp_count = NULL;
    }
    free(thread_pt);
    thread_pt = NULL;
    *bcsc_struct_pointer = bcolptr;
    *bcsc_struct_index = browidx;
    *bcsc_struct_nnzptr = bnnzptr;
    
    *bcsc_inblock_pointers = malloc(sizeof(analyzer_inblock_ptr*) * block_num);
    *bcsc_inblock_indeces = malloc(sizeof(analyzer_inblock_idx*) * block_num);
    if(csc_value){
        *bcsc_values = malloc(sizeof(sparse_value_t*) * block_num);
    }

    for (analyzer_exblock_idx bcol = 0; bcol < block_length; bcol++)
    {
        for (analyzer_exblock_ptr bidx = bcolptr[bcol]; bidx < bcolptr[bcol+1]; bidx++)
        {
            analyzer_inblock_ptr *cur_block_colptr = (analyzer_inblock_ptr *)(block_csc + sizeof(analyzer_inblock_ptr) * bidx * (nb + 1));
            analyzer_inblock_idx *cur_block_rowidx = (analyzer_inblock_idx *)(
                block_csc
                + sizeof(analyzer_inblock_ptr) * bcolptr[block_length] * (nb + 1) 
                + sizeof(analyzer_inblock_idx) * bnnzptr[bidx]
            );
            sparse_value_t *cur_block_value = NULL;
            if(csc_value){
                cur_block_value = (sparse_value_t *)(
                    block_csc
                    + sizeof(analyzer_inblock_ptr) * bcolptr[block_length] * (nb + 1) 
                    + sizeof(analyzer_inblock_idx) * bnnzptr[bcolptr[block_length]]
                    + sizeof(sparse_value_t) * bnnzptr[bidx]
                );
            }

            (*bcsc_inblock_pointers)[bidx] = cur_block_colptr;
            (*bcsc_inblock_indeces)[bidx] = cur_block_rowidx;
            if(csc_value){
                (*bcsc_values)[bidx] = cur_block_value;
            }
        }
    }

#undef _PANGULU_SET_BVALUE_SIZE
}

int analyzer_structed_sparse(sparse_matrix_t* s, int nb){
    analyzer_exblock_ptr* bcsc_struct_pointer;
    analyzer_exblock_idx* bcsc_struct_index;
    analyzer_exblock_ptr* bcsc_struct_nnzptr;
    analyzer_inblock_ptr** bcsc_inblock_pointers;
    analyzer_inblock_idx** bcsc_inblock_indeces;
    sparse_value_t** bcsc_value;

    printf("nb = %d\n", nb);

    analyzer_csc_to_bcsc(
        s->n,
        nb,

        s->rowpointer,
        s->columnindex,
        s->value,

        &bcsc_struct_pointer,
        &bcsc_struct_index,
        &bcsc_struct_nnzptr,
        &bcsc_inblock_pointers,
        &bcsc_inblock_indeces,
        &bcsc_value
    );

    analyzer_exblock_idx n = s->n;
    analyzer_int64_t block_length = ICEIL(n, nb);

    analyzer_int64_t max_dist_sum = 0;
    analyzer_int64_t nnz_structed = 0;

    for(analyzer_exblock_idx bcol = 0; bcol < block_length; bcol++){
        for(analyzer_exblock_idx bidx=bcsc_struct_pointer[bcol];bidx<bcsc_struct_pointer[bcol+1];bidx++){
            analyzer_exblock_idx brow = bcsc_struct_index[bidx];
            analyzer_inblock_ptr* colpointer = bcsc_inblock_pointers[bidx];
            analyzer_inblock_idx* rowindex = bcsc_inblock_indeces[bidx];
            sparse_value_t* value = bcsc_value[bidx];
            analyzer_int32_t max_dist = 0; 
            // printf("((%d, %d)) %d\n", brow, bcol, colpointer[nb]);
            for(analyzer_inblock_idx col = 0; col < nb; col++){
                for(analyzer_inblock_ptr idx = colpointer[col]; idx < colpointer[col+1]; idx++){
                    analyzer_inblock_idx row = rowindex[idx];
                    // printf("(%d, %d) ", row, col);
                    analyzer_int32_t dist = abs(row - col);
                    if(dist>max_dist){
                        max_dist = dist;
                    }
                }
            }
            // printf("%d\n", max_dist);
            max_dist_sum+=max_dist;
            nnz_structed += ((2 * max_dist + 1) * nb);
        }
    }
    
    printf("Avg. max_dist : %.2lf\n", ((double)max_dist_sum) / bcsc_struct_pointer[block_length]);
    printf("nnz after filling : %ld rate : %.2lfx\n", nnz_structed, ((double)nnz_structed) / s->rowpointer[n]);

    // printf("block count : %d\n", bcsc_struct_pointer[block_length]);

    // // int ncol = 0;
    // int nnz = 0;
    // for(long long bidx=0;bidx<bcsc_struct_pointer[block_length];bidx++){
    //     for(int col = 0; col < nb; col++){
    //         int nzcol = bcsc_inblock_pointers[bidx][col+1] - bcsc_inblock_pointers[bidx][col];
    //         nnz+=nzcol;
    //     }
    // }
    // printf("nnz from blocks = %d\n", nnz);

}